import torch
import numpy as np
import os

class AuxInDistributionImages(torch.utils.data.Dataset):

    def __init__(self, aux_file='../../datasets/aux_indistribution/cifa10-1m.npz', transform=None):
        assert os.path.exists(aux_file)
        self.transform = transform
        data = np.load(aux_file)
        self.image = data['image'] #.astype(np.float16)
        self.label = data['label']
        print(f"Auxillary In-distribution Images - Length {len(data['image'])}")

    def __getitem__(self, index):
        x = self.image[index]
        y = self.label[index]
        if self.transform is not None:
            x = self.transform(x)
        return x, y

    def __len__(self):
        return len(self.image)


class AuxTiny200Images(torch.utils.data.Dataset):

    def __init__(self, aux_path='../../datasets/aux_indistribution/', file_x='tiny_edm_1m_x.npy',
                 file_y='tiny_edm_1m_y.npy', length=1000000, img_size=64, transform=None):
        aux_file_x = os.path.join(aux_path, file_x)
        aux_file_y = os.path.join(aux_path, file_y)
        assert os.path.exists(aux_file_x) and os.path.exists(aux_file_y)
        self.transform = transform
        self.x_shape = (length, img_size, img_size, 3)
        self.y_shape = (length,)
        self.memapx = np.memmap(aux_file_x, mode='r', dtype='uint8', shape=self.x_shape)
        self.memapy = np.memmap(aux_file_y, mode='r', dtype='uint8', shape=self.y_shape)
        print(f"Auxillary In-distribution Images - Length {len(self.memapx)}")

    def __getitem__(self, index):
        x = self.memapx[index]
        y = self.memapy[index]
        x = np.array(x)
        if self.transform is not None:
            x = self.transform(x)
        return x, y

    def __len__(self):
        return self.y_shape[0]
